import json
from rank_bm25 import BM25Okapi
import jieba

def get_bm25_similarity_top10(query, questions_data):
    """
    使用rank_bm25库获取与查询最相似的前10个问题ID
    :param query: 查询字符串
    :param questions_data: 问题数据列表
    :return: 前10个最相似问题的ID列表
    """
    # 提取问题文本列表
    corpus = [q['question'] for q in questions_data]
    
    # 对语料库进行分词
    tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus]
    
    # 初始化BM25模型
    bm25 = BM25Okapi(tokenized_corpus)
    
    # 对查询进行分词
    tokenized_query = list(jieba.cut(query))
    
    # 计算得分
    scores = bm25.get_scores(tokenized_query)
    
    # 获取得分最高的前10个文档的索引
    top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:10]
    
    # 根据索引获取对应的ID
    top_ids = [questions_data[i]['id'] for i in top_indices]
    
    return top_ids


processed_id = []
group_id = []
while len(processed_id)<200:

    # 读取问题数据
    questions_data = []
    with open('question_c.json', encoding='utf8') as f:
        for line in f:
            data = json.loads(line)
            if data['id'] not in processed_id:
                questions_data.append(data)

    # 示例用法：
    query = questions_data[0]['question']
    top10_ids = get_bm25_similarity_top10(query, questions_data)

    doc = ''
    for i in questions_data[1:]:
        if i['id'] in top10_ids:
            doc+=json.dumps(i,indent=4,ensure_ascii=False)
            doc+='\n---\n'

    prompt = f'''
    已知原题目为:{query}
    请找到以下题目和原题目解法步骤相同的题目，并输出ID
    题目为:
    {doc}
    
    解法和步骤相同的题目的含义为：题目已知信息类型完全一致，结果参数完全一致，示例：
    一致的题目：
    我想知道长春高榕生物科技有限公司作为被告的案件中涉案金额小于100万大于1万的案号分别为？涉案金额数值为？
    伊吾广汇矿业有限公司作为被告的案件中涉案金额小于10万大于1万的案号分别为？涉案金额数值为？
    不一致的题目：
    我想知道长春高榕生物科技有限公司作为被告的案件中涉案金额小于100万大于1万的案号分别为？涉案金额数值为？
    伊吾广汇矿业有限公司作为原告的案件中涉案金额小于100万大于1万的案号分别为？
    不一致原因：一个是作为被告一个是作为原告，参数不一样，另外涉案金额数值没有，需求结果不一样

    按照以下json格式输出:
    '''+'''
    ```json
    {"same_ids":[int,...]}
    ```
    '''

    from llm import super_eval,llm
    from copy import deepcopy
    print(prompt)
    res = super_eval(llm(prompt))
    processed_id.extend(res['same_ids'])
    processed_id.append(questions_data[0]['id'])
    group = deepcopy(res['same_ids'])
    group.append(questions_data[0]['id'])
    group_id.append(group)
    # 保存group_id
    with open('group_id.json', 'w', encoding='utf8') as f:
        json.dump(group_id, f, ensure_ascii=False, indent=4)